#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
对Exx系列进行外插处理
"""

import argparse
import os
import sys
import numpy as np
from pylab import plt
import pandas as pd
from obspy.core.utcdatetime import UTCDateTime
import glob

# 导入自定义工具
from utils.h5data import read_h5_data
from utils.trace import filtfilt
from utils.plot import plot_traces, plot_raw_data
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.math import norm

# 命令行参数解析
parser = argparse.ArgumentParser(description='Plot traceSeq h5 data from multiple files')
parser.add_argument('-input', default='', help='input h5 file pattern')
parser.add_argument('-figroot', default='figures/debug', help='root to save figs')
parser.add_argument('-noD1', action='store_true', help='method: discard S2N events')
parser.add_argument('-noD2', action='store_true', help='method: discard N2S events')
parser.add_argument('-noD3', action='store_true', help='method: discard mixed events')
args = parser.parse_args()


INPUT_PATTERN = args.input
FIG_ROOT = args.figroot
USE_D1_EVENTS = not args.noD1
USE_D2_EVENTS = not args.noD2
USE_D3_EVENTS = not args.noD3

# 获取所有匹配的文件
if not INPUT_PATTERN:
    print("请提供文件匹配模式，例如: data/2406.HSR/events/P192.2406.Z.traces.P*.h5")
    sys.exit(1)

input_files = glob.glob(INPUT_PATTERN)
input_files.sort()
if not input_files:
    print(f"未找到匹配的文件: {INPUT_PATTERN}")
    sys.exit(1)

print(f"找到 {len(input_files)} 个文件")

# 读取第一个文件获取基本参数
datasets, args_infile = read_h5_data(input_files[0], keys=['MARKER','EMARKER'], group_name='metadata', read_attrs=True)
DATE = args_infile['date']
fs = args_infile['fs']
fe = args_infile['fe']
EMARKER = datasets[1].decode()
MARKER = datasets[0].decode()
print(args_infile)
FIG_ROOT = f'{FIG_ROOT}/{DATE}.6E.all'
if not os.path.exists(FIG_ROOT):
    os.makedirs(FIG_ROOT)

# 存储所有数据
all_ref_traces = []
all_x = []
all_corr_stats = []
train_pos = []
# 读取所有文件的数据
print('Reading data from all files...')
for i, input_file in enumerate(input_files):
    print(f'  Reading {i+1}/{len(input_files)}: {os.path.basename(input_file)}')
    
    # 读取h5文件中的数据
    try:
        data_dict = read_h5_data(input_file, keys=['t','ref', 'S', 'R', 'x'])
        _, args_infile1 = read_h5_data(input_file, keys=['MARKER'], group_name='metadata', read_attrs=True)
        tn, ref_traces, S_data, R, x = data_dict
        vinfo = read_h5_data(input_file, keys=['info'])[0]
        # 解码字符串数组
        S = S_data.decode()
        R = [R[i].decode() for i in range(len(R))]
        ns,nt = ref_traces.shape
        # idx_events = np.where(vinfo==2)[0]
        
        r_shift = int(S[1:])*3-3
        # 校正X
        x = x-x[R.index('P218')]-r_shift

        # R,S 转换
        x = -x
        ref_traces = ref_traces[:,-1::-1]

        
        print(S,r_shift)
                              
        for j in range(ns):
            name_j = R[j]
            if 'E' not in name_j :
                all_corr_stats.append(S+'/'+name_j)
                all_ref_traces.append(ref_traces[j,:])
                # all_ref_traces.append(traces[idx_events,j,:].sum(axis=0))
                all_x.append(x[j])
                # print(S+'/'+name_j,x[j])
        train_pos.append(args_infile1['Tpos'])
            
    except Exception as e:
        print(f"  跳过文件 {input_file}，错误: {e}")

# 合并所有数据
if not all_ref_traces:
    print("没有成功读取任何数据")
    sys.exit(1)

ref_traces_combined = np.array(all_ref_traces)
x_combined = np.array(all_x)
# 按距离排序
sort_idx = np.argsort(x_combined)
x_combined = x_combined[sort_idx]
ref_traces_combined = ref_traces_combined[sort_idx,:]
all_corr_stats = [all_corr_stats[i] for i in sort_idx]

ns = len(all_corr_stats)

print(f'总共有 {ns} 个 traces')

# 数据处理
ref_traces_combined = filtfilt(ref_traces_combined, tn[1]-tn[0], [fs, fe], order=4)
ref_traces_combined = norm(ref_traces_combined, ONE_AXIS=True)

# 转换为numpy数组
ref_traces = np.array(ref_traces_combined)
x = np.array(x_combined)

print(ref_traces_combined.shape)

# 绘制参考道数据
print('Plotting reference traces...')

# 参数设置
Tpos='N'*('N' in train_pos) + 'S'*('S' in train_pos)
print(Tpos)

fs, fe = 5, 8
VLIM = [100, 4000]
TLIM = [-3, 3]

# fs, fe = 1, 4
# VLIM = [100, 1000]
# TLIM = [-5, 5]

ref_traces = filtfilt(ref_traces, tn[1]-tn[0], [fs, fe], order=4)
ref_traces = norm(ref_traces, ONE_AXIS=True)

NORM = True
PLOT_WIGGLE = True if ns < 30 else False
PLOT_WIGGLE = False

# 绘制图形
fig, ax1, ax2 = plot_raw_data(ref_traces, x, tn, fs=fs, fe=fe, 
                              VLIM=VLIM, PLOT_WIGGLE=PLOT_WIGGLE, SCALE=6, FV_NORM=NORM)

if PLOT_WIGGLE:
    # for j in range(30):  # 限制标签数量避免过于拥挤
    #     ax1.text(-1, x[j], all_corr_stats[j], fontsize=6)
    ax1.set_ylim([x.min()-100, x.max()+100])
else:
    ax1.set_ylim([x.min(), x.max()])

ax1.set_title(f'ALL.F{fs:.1f}.{fe:.1f}.NORM{int(NORM)}.Tp.{Tpos}')
ax1.set_xlim(TLIM)

# 频散曲线
try:
    disp = np.loadtxt('figures/6.traceSeq.figures/disp.txt', delimiter=',')
    ax2.plot(disp[:,0], disp[:,1], 'yellow', lw=1)
    ax2.plot(disp[:,0], -disp[:,1], 'yellow', lw=1)
except:
    print("未找到 disp.txt 文件，跳过频散曲线绘制")

# # 参考速度
# ax1.plot([-5, 5], [-5*3000, 5*3000], 'k-', lw=1)
# ax1.plot([-5, 5], [-5*340, 5*340], 'k--', lw=1)

# # 添加特定频率线
# f0 = 4.4
# for j in [1, 4, 9]:
#     ax2.plot([j*f0, j*f0], VLIM, color='yellow', lw=1)

ax2.set_ylim(VLIM)
fig.tight_layout()

# 保存图像
figname = f'{FIG_ROOT}/ref.{fs:03.1f}.{fe:03.1f}.NORM{int(NORM)}.combined.Tp.{Tpos}.png'
fig.savefig(figname, dpi=300)
print(f'Saved {figname}')

print('Plotting completed.')